#!/usr/bin/python

import Image, math, os, sys
from ImageDraw import *
from numpy import *
from opencv.adaptors import *

MOSAIC_SHAPE_SQUARES=0
MOSAIC_SHAPE_CIRCLES=1
MOSAIC_SHAPE_POLYGON=2

class Mosaic:
    def __init__(self, filename, mapsize=[2048,2048], indexed = False, useAvg = False, bw=False):
        self.numhigh = 0
        self.bw = bw
        self.indexed = indexed
        self.useAvg = useAvg
        self.filename = filename
        self.mapsize = array(mapsize)
        self.mapim = Image.new('RGBA', mapsize, (0, 0, 0, 0))
        self.big_im = None
        self.shape = 0
        self.sides = 4
        self.reset()
    
    def setBW(self, bw):
        self.bw = bw
        self.reset()

    def setShape(self, shape, sides=4):
        self.shape = shape
        self.sides = sides
        self.reset()
        
    def reset(self):
        if self.bw:
            self.nims = array([16,16])
            self.min = 255; self.max = 0
        elif self.indexed:
            self.nims = array([16,16])
            self.rmin = self.gmin = self.bmin = ones((6,6))*255
            self.rmax = self.gmax = self.bmax = zeros((6,6))
        else:
            self.nims = array([64, 64])
            self.rmin = self.gmin = self.bmin = ones((16,16))*255
            self.rmax = self.gmax = self.bmax = zeros((16,16))

        self.submapsize = self.mapsize/self.nims
        if self.shape == MOSAIC_SHAPE_SQUARES:
            self.fromSquares()
        if self.shape == MOSAIC_SHAPE_CIRCLES:
            self.fromCircles()
        if self.shape == MOSAIC_SHAPE_POLYGON:
            self.fromPolygon(self.sides)

    def mapAvg(self, r, b, g):        
        if self.indexed:
            ri = int(floor(r/51)); gi = int(floor(g/51)); bi = int(floor(b/51))

        else:
            ri = int(floor(r/16)); gi = int(floor(g/16)); bi = int(floor(b/16))


        if self.bw:
#            rmin = 0; rmax = 75  
#            gmin = 0; gmax = 75  
#            bmin = 0; bmax = 75  
            rmin = 40; rmax = 125  
            gmin = 40; gmax = 125  
            bmin = 40; bmax = 125  
        else:
            rmin = self.rmin[gi,bi]; rmax = self.rmax[gi,bi]  
            gmin = self.gmin[ri,bi]; gmax = self.gmax[ri,bi]  
            bmin = self.bmin[ri,gi]; bmax = self.bmax[ri,gi]  
            
        if (rmax-rmin) > 0:        
            r=(r-rmin)/float(rmax-rmin)*255
        else:
            r=0
        if (gmax-gmin) > 0:        
            g=(g-gmin)/float(gmax-gmin)*255
        else:
            g=0
        if (bmax-bmin) > 0:        
            b=(b-bmin)/float(bmax-bmin)*255
        else:
            b=0

        return r,g,b

    def addImage(self, subImg):
        
        #subImg = subImg.resize(self.submapsize, Image.BICUBIC)
        subImg = subImg.resize(self.submapsize)
        nImg = PIL2NumPy(subImg)
        pixels = self.submapsize[0]*self.submapsize[1]

        red = nImg[:,:,0]
        green = nImg[:,:,1]
        blue = nImg[:,:,2]
        r = float(red.sum())/pixels
        g = float(green.sum())/pixels
        b = float(blue.sum())/pixels

        if self.useAvg:
            self.getAvg(subImg)

        if self.bw:
            r,g,b = self.mapAvg(r, g, b)
            color = int((r + g + b)/3.)
            #print color
            if color > 255:
                return
        elif self.indexed:
            r = int(floor(r/51)); g = int(floor(g/51)); b = int(floor(b/51))
            color = r + 6*g + 6*6*b
        else:
            r = int(floor(r/16)); g = int(floor(g/16)); b = int(floor(b/16))
            bx = b%4; by = b/4
            color = r + bx*16 + g*64 + by*1024


        color = max(color, 0)
        color = min(self.nims[0]*self.nims[1]-1, color)

#        if color > 250:
#            subImg.save("snowflake_high/%04d.jpg"%(self.numhigh))
#            self.numhigh = self.numhigh+1

        x = color%self.nims[0]
        y = color/self.nims[1]

        box = (array([[x, y], [x+1, y+1]]) * self.submapsize).reshape(4)
        self.mapim.paste(subImg, tuple(box))

    def getAvg(self, img):
        img = img.convert("RGB").resize(self.submapsize)
        nImg = PIL2NumPy(img)
        pixels = self.submapsize[0]*self.submapsize[1]

        red = nImg[:,:,0]
        green = nImg[:,:,1]
        blue = nImg[:,:,2]

        r = float(red.sum())/pixels
        g = float(green.sum())/pixels
        b = float(blue.sum())/pixels

        if self.indexed:
            ri = int(floor(r/51)); gi = int(floor(g/51)); bi = int(floor(b/51))

        else:
            ri = int(floor(r/16)); gi = int(floor(g/16)); bi = int(floor(b/16))

        if r > self.rmax[gi,bi]:
            self.rmax[gi,bi] = r        
        if r < self.rmin[gi,bi]:
            self.rmin[gi,bi] = r        
        if g > self.gmax[ri,bi]:
            self.gmax[ri,bi] = g        
        if g < self.gmin[ri,bi]:
            self.gmin[ri,bi] = g        
        if b > self.bmax[ri,gi]:
            self.bmax[ri,gi] = b        
        if b < self.bmin[ri,gi]:
            self.bmin[ri,gi] = b        

    def transformDir(self, dir, imsize, nbig_pix):        
        files = os.listdir(dir)
        for file in files:
            if not file == ".DS_Store":
                filename = dir+"/"+file
                print "mosaicing " + filename
                self.makeMosaic(dir+"/mosaic_"+file, Image.open(filename).convert("RGBA"), imsize, nbig_pix)

    def fromDir(self, dir):
        files = os.listdir(dir)

        for file in files:
            print "adding " + dir+"/"+file

            subImg = Image.open(dir+"/"+file)
            if self.bw:
                subImg = subImg.convert("L").convert("RGB").resize(self.submapsize, Image.BICUBIC)

            print self.submapsize
            self.addImage(subImg)
            
            # this code adds scales the subimages to a bunch of different sizes 
            # and adds them to the mosaic map to make it more dense
            img = Image.new('RGB', (128,128))
            img.paste(subImg.resize((112,112), Image.BICUBIC), (8,8,120,120))
            self.addImage(img)
            img = Image.new('RGB', (128,128))
            img.paste(subImg.resize((96,96), Image.BICUBIC), (16,16,112,112))
            self.addImage(img)
            img = Image.new('RGB', (128,128))
            img.paste(subImg.resize((80,80), Image.BICUBIC), (24,24,104,104))
            self.addImage(img)
            img = Image.new('RGB', (128,128))
            img.paste(subImg.resize((64,64), Image.BICUBIC), (32,32,96,96))
            self.addImage(img)
            img = Image.new('RGB', (128,128))
            img.paste(subImg.resize((48,48), Image.BICUBIC), (40,40,88,88))
            self.addImage(img)
            img = Image.new('RGB', (128,128))
            img.paste(subImg.resize((32,32), Image.BICUBIC), (48,48,80,80))
            self.addImage(img)

        self.mapim.save(self.filename)

    def makeMosaic(self, img, imsize, nbig_pix):

        alpha = .0
        imsize = array(imsize)
        nbig_pix = array(nbig_pix)
        pix_imsize = imsize/nbig_pix

#        if self.bw:
#            img.convert("L").convert("RGBA")

#        if not self.big_im:
#            self.big_im = Image.new('L', imsize, 0)

        if self.bw:
            self.big_im = Image.new('L', imsize, 0)
        else:
            self.big_im = Image.new('RGBA', imsize, 0)

        #big_im = Image.new('RGBA', imsize, (0, 0, 0, 0))
        img = img.resize(nbig_pix)

        for i in range(nbig_pix[0]):
            for j in range(nbig_pix[1]):

                v = img.getpixel((i,j))
                r = v[0]; g = v[1]; b = v[2]

                if self.useAvg:
                    r,g,b = self.mapAvg(r, g, b)

                    if self.indexed:
                        r = int(floor(r/51)); g = int(floor(g/51)); b = int(floor(b/51))
                        val = r + 6*g + 6*6*b
                    else:
                        r = int(floor(r/16)); g = int(floor(g/16)); b = int(floor(b/16))
                        bx = b%4; by = b/4
                        val = r + bx*16 + g*64 + by*1024
                else:
                    if self.bw:
                        val = (r + g + b)/3  
                    elif self.indexed:
                        r = r/51; g = g/51; b = b/51
                        val = r + 6*g + 6*6*b  

                    else:
                        r = r/16; g = g/16; b = b/16
                        bx = b%4; by = b/4
                        val = r + bx*16 + g*64 + by*1024  
                    
                x = val%self.nims[0]
                y = val/self.nims[1]

                # get the mosaic sub image for this pixel value
                mapbox = tuple((array([[x, y], [x+1, y+1]]) * self.submapsize).reshape(4))
                #littleim = self.mapim.crop(mapbox).resize(pix_imsize, Image.BICUBIC)
                littleim = self.mapim.crop(mapbox).resize(pix_imsize)
                            
                # fill the final image with the mosaic sub-image we just made
                bigbox = tuple((array([[i, j], [i+1, j+1]]) * array(pix_imsize)).reshape(4))
                self.big_im.paste(littleim, bigbox)

        return self.big_im.convert('RGBA')

    def fromPolygon(self, sides):
        if self.bw:
            for x in range(self.nims[0]):
                for y in range(self.nims[1]):
                    c = (x + y*16)/256.0                            
                    box = tuple((array([[x, y], [x+1, y+1]]) * self.submapsize).reshape(4))

                    sbox = []
                    dtheta = 2*pi/sides
                    for i in range(sides):
                        xx = (.5*cos(i*dtheta + pi/4)+.5)*self.submapsize[0]
                        yy = (.5*sin(i*dtheta + pi/4)+.5)*self.submapsize[1]
                        sbox.append([xx, yy])
                    sbox = array(sbox)
                    #sbox = array([[0, 0], [self.submapsize[0], 0], [self.submapsize[0]/2., self.submapsize[1]]])
                    cen = array(self.submapsize)/2

                    i = Image.new("L", tuple(self.submapsize))
                    d = Draw(i)
                    d.polygon(tuple((c*sbox + (1-c)*cen).reshape(2*sides,1)), 255)
                    self.mapim.paste(i, tuple(box))                    

        elif not self.indexed:    
            for by in range(4):
                for g in range(16):
                    for bx in range(4):
                        for r in range(16):
                            b = bx + 4*by
                            color = r + bx*16 + g*64 + by*1024
                            x = color%self.nims[0]
                            y = color/self.nims[1]
                            
                            box = tuple((array([[x, y], [x+1, y+1]]) * self.submapsize).reshape(4))

                            sbox = []
                            dtheta = 2*pi/sides
                            for i in range(sides):
                                xx = (.5*cos(i*dtheta + pi/4)+.5)*self.submapsize[0]
                                yy = (.5*sin(i*dtheta + pi/4)+.5)*self.submapsize[1]
                                sbox.append([xx, yy])
                            sbox = array(sbox)
                            #sbox = array([[0, 0], self.submapsize])
                            cen = array(self.submapsize)/2

                            rr = r/16.; gg = g/16.; bb = b/16.
                            ri = Image.new("L", tuple(self.submapsize))
                            gi = Image.new("L", tuple(self.submapsize))
                            bi = Image.new("L", tuple(self.submapsize))
                            rd = Draw(ri)
                            gd = Draw(gi)
                            bd = Draw(bi)
                            rd.ellipse(tuple((rr*sbox + (1-rr)*cen).reshape(2*sides,1)), 255)
                            gd.ellipse(tuple((gg*sbox + (1-gg)*cen).reshape(2*sides,1)), 255)
                            bd.ellipse(tuple((bb*sbox + (1-bb)*cen).reshape(2*sides,1)), 255)
                            i = Image.merge("RGB", (ri, gi, bi))
                            self.mapim.paste(i, tuple(box))                    


    def fromSquares(self):
        if self.bw:
            for i in range(self.nims[0]):
                for j in range(self.nims[1]):
                    c = i + j*16                            
                    valI = Image.new("RGBA", self.submapsize, (c, c, c, 1)) 
                    # fill the final image with the mosaic sub-image we just made
                    bigbox = tuple((array([[i, j], [i+1, j+1]]) * array(self.submapsize)).reshape(4))
                    self.mapim.paste(valI, bigbox)
        elif not self.indexed:
            for by in range(4):
                for g in range(16):
                    for bx in range(4):
                        for r in range(16):
                            b = bx + 4*by
                            color = r + bx*16 + g*64 + by*1024
                            i = color%self.nims[0]
                            j = color/self.nims[1]
                            rr = r*16.; gg = g*16.; bb = b*16.
                            valI = Image.new("RGBA", self.submapsize, (rr, gg, bb, 1)) 
                            # fill the final image with the mosaic sub-image we just made
                            bigbox = tuple((array([[i, j], [i+1, j+1]]) * array(self.submapsize)).reshape(4))
                            self.mapim.paste(valI, bigbox)

    def fromCircles(self):
        if self.bw:
            for x in range(self.nims[0]):
                for y in range(self.nims[1]):
                    c = (x + y*16)/256.0                            
                    box = tuple((array([[x, y], [x+1, y+1]]) * self.submapsize).reshape(4))

                    sbox = array([[0, 0], self.submapsize])
                    cen = array(self.submapsize)/2

                    i = Image.new("L", tuple(self.submapsize))
                    d = Draw(i)
                    d.ellipse(tuple((c*sbox + (1-c)*cen).reshape(4,1)), 255)
                    self.mapim.paste(i, tuple(box))                    

        elif not self.indexed:    
            for by in range(4):
                for g in range(16):
                    for bx in range(4):
                        for r in range(16):
                            b = bx + 4*by
                            color = r + bx*16 + g*64 + by*1024
                            x = color%self.nims[0]
                            y = color/self.nims[1]
                            
                            box = tuple((array([[x, y], [x+1, y+1]]) * self.submapsize).reshape(4))

                            sbox = array([[0, 0], self.submapsize])
                            cen = array(self.submapsize)/2

                            rr = r/16.; gg = g/16.; bb = b/16.
                            ri = Image.new("L", tuple(self.submapsize))
                            gi = Image.new("L", tuple(self.submapsize))
                            bi = Image.new("L", tuple(self.submapsize))
                            rd = Draw(ri)
                            gd = Draw(gi)
                            bd = Draw(bi)
                            rd.ellipse(tuple((rr*sbox + (1-rr)*cen).reshape(4,1)), 255)
                            gd.ellipse(tuple((gg*sbox + (1-gg)*cen).reshape(4,1)), 255)
                            bd.ellipse(tuple((bb*sbox + (1-bb)*cen).reshape(4,1)), 255)
                            i = Image.merge("RGB", (ri, gi, bi))
                            self.mapim.paste(i, tuple(box))                    
        
    
def main( *argv):
    # initialize a mosaic image look-up table
    m = Mosaic("snowflake.jpg", [2048,2048], indexed=False, useAvg=False, bw=True)
    # fill the look-up table with images from a directory
    #m.fromDir("snowflake")
    m.mapim = Image.open("snowflake.jpg").convert("RGBA")
    # open a photo and turn it into a photomosaic form the map you created
    toMosaic = Image.open("stephollie.jpg").convert("RGBA")
    m.makeMosaic(toMosaic, [2*1536, 2*2048], [36, 48]).save("snowflake_mosaic.jpg")    

if __name__ == '__main__':
    sys.exit( main( *sys.argv))

        
